#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import os
import sys


class ThrRole(object):
    class _ThrRole(object):
        def __init__(self, id, name, feature):
            self.id = id
            self.name = name
            self.feature = feature

    UNKNOWN      = _ThrRole(-1, "unknown",         "")
    PM           = _ThrRole(0,  "postmaster",      "PostmasterMain")
    WORKER       = _ThrRole(1,  "worker",          "GaussDbThreadMain<(knl_thread_role)1>")
    TPLWORKER    = _ThrRole(2,  "TPLworker",       "GaussDbThreadMain<(knl_thread_role)2>")
    TPLLISTEN    = _ThrRole(3,  "TPLlistener",     "TpoolListenerMain")
    TPLSCHEDUL   = _ThrRole(4,  "TPLscheduler",    "TpoolSchedulerMain")
    AVCLAUNCHER  = _ThrRole(7,  "AVClauncher",     "AutoVacLauncherMain")
    AVCWORKER    = _ThrRole(8,  "AVCworker",       "AutoVacWorkerMain")
    JOB          = _ThrRole(9,  "Jobscheduler",    "JobScheduleMain")
    WLMWORKER    = _ThrRole(11, "WLMworker",       "WLMProcessThreadMain")
    WLMMONITOR   = _ThrRole(12, "WLMmonitor",      "WLMmonitorMain")
    WLMARBITER   = _ThrRole(13, "WLMarbiter",      "WLMarbiterMain")
    AUDITOR      = _ThrRole(15, "auditor",         "PgAuditorMain")
    STATS        = _ThrRole(16, "statscollector",  "PgstatCollectorMain")
    SYSLOGGER    = _ThrRole(17, "syslogger",       "SysLoggerMain")
    ALARM        = _ThrRole(20, "alarm",           "AlarmCheckerMain")
    REAPER       = _ThrRole(21, "reaper",          "ReaperBackendMain")
    TWOPCCLEANER = _ThrRole(23, "2pccleaner",      "TwoPhaseCleanerMain")
    FAULTMONITOR = _ThrRole(25, "faultmonitor",    "FaultMonitorMain")
    SPBGWRITER   = _ThrRole(27, "Spbgwriter",      "invalid_buffer_bgwriter_main")
    PERCENTILE   = _ThrRole(28, "percentworker",   "PercentileMain")
    TXNSNAPCAP   = _ThrRole(29, "txnsnapcapturer", "TxnSnapCapturerMain")
    CFSSHRINKER  = _ThrRole(31, "CfsShrinker",     "CfsShrinkerMain")
    ASHWORKER    = _ThrRole(35, "ashworker",       "ActiveSessionCollectMain")
    TRACKSTMT    = _ThrRole(36, "TrackStmtWorker", "StatementFlushMain")
    CKPT         = _ThrRole(38, "checkpointer",    "CheckpointerMain")
    WALWRITER    = _ThrRole(39, "WALwriter",       "WalWriterMain")
    WALWRITERAUX = _ThrRole(40, "WALwriteraux",    "WalWriterAuxiliaryMain")
    PAGEWRITER   = _ThrRole(46, "pagewriter",      "ckpt_pagewriter_main")
    UNDOLAUNCH   = _ThrRole(57, "asyncundolaunch", "UndoLauncherMain")
    UNDORECYCLER = _ThrRole(56, "undorecycler",    "UndoRecycleMain")
    GLOBALSTATS  = _ThrRole(60, "globalstats",     "GlobalStatsTrackerMain")
    APPLYLAUNCH  = _ThrRole(65, "applylauncher",   "ApplyLauncherMain")
    MESWORKER    = _ThrRole(1001, "mes_worker",      "mes_task_proc")
    MESRECV      = _ThrRole(1002, "mes_recv",        "mes_recv_proc")
    MESHEART     = _ThrRole(1003, "mes_heartbeat",   "mes_heartbeat_entry")
    DMSPARALLEL  = _ThrRole(1004, "dms_parallel",    "dms_reform_parallel_thread")
    DMSAUX       = _ThrRole(1005, "dms_auxiliary",   "DmsAuxiliaryMain")

    @staticmethod
    def identify(bt):
        for k, role in ThrRole.__dict__.items():
            if type(role) != ThrRole._ThrRole:
                continue
            if k == "UNKNOWN":
                continue
            if bt.find(role.feature) >= 0:
                return role

        return ThrRole.UNKNOWN

    @staticmethod
    def get_role_map():
        res = []
        for k, role in ThrRole.__dict__.items():
            if type(role) != ThrRole._ThrRole:
                continue
            res.append(f"[{role.id} {role.name}]")
        return '\t'.join(res)


class Thread(object):
    def __init__(self, bt):
        lines = bt.splitlines()
        self.src = bt

        self.id = lines[0].split()[1]
        self.role = ThrRole.identify(bt)
        self.tid = lines[0].split("LWP ")[1].split(')')[0]

        self.title = lines[0]
        self.backtraces = lines[1:]
        self.functraces = []
        self.functraces_func = []
        for line in self.backtraces:
            parts = line.split()
            functrace = f"{parts[0]} {parts[3]}(...)" if parts[1].startswith("0x") else f"{parts[0]} {parts[1]}(...)"
            self.functraces.append(functrace)
            self.functraces_func.append(functrace.split()[1].split('(')[0])


class ThreadList(object):
    """
    保存一个线程列表
    """
    def __init__(self):
        self.threads = []

    def append(self, thread):
        self.threads.append(thread)

    def show(self, detail=True, summary=True):
        if detail:
            for thread in self.threads:
                print(thread.src, "\n")

        if summary:
            ids = self.ids()
            print("线程总数:", len(ids))
            print("线程id列表:", ' '.join(ids))

    def ids(self):
        return [thread.id for thread in self.threads]


class ThreadGroup(object):
    """
    保存线程分组
    """
    
    def __init__(self, group_func):
        self._group_func = group_func
        self.group = {}
        self.group_key = []

    def append(self, key, thread):
        if self.group.get(key) is None:
            self.group[key] = ThreadList()
            self.group_key.append(key)
        self.group[key].append(thread)

    def show(self, detail=True, summary=True):
        all_thread_count = 0
        for i, key in enumerate(self.group_key):
            all_thread_count += len(self.group[key].ids())
            if not detail:
                continue
            print(f"[组 {i} ]:——————————————————————————————————————————————————————")
            print(key)
            self.group[key].show(detail=False, summary=True)
            print()

        if summary:
            print("线程总数:", all_thread_count)
            print("分组总数:", len(self.group))

    def open(self, i):
        Assert(i < len(self.group_key), f"没有第 {i} 组")
        key = self.group_key[i]
        return self.group[key]


class Operate(object):
    def __init__(self, name, param_count, input_type, support_reverse, usage):
        self.name = name
        self.param_count = param_count
        self.input_type = input_type
        self.support_reverse = support_reverse
        self.usage = usage

    def execute(self, thread_list, params, reverse=False):
        assert False


def Assert(condition, msg=None):
    if condition:
        return
    if msg is not None:
        print(msg)
    exit(1)


###############################
# 解析文件或者执行“pstack -p”抓取
###############################
class Pstack(object):
    @staticmethod
    def parse(target):
        if os.access(target, os.F_OK):
            with open(target, 'r') as f:
                lines = f.readlines()
        else:
            import subprocess
            Assert(target.isdigit(), f"找不到文件{target}")
            status, output = subprocess.getstatusoutput(f"pstack {target}")
            Assert(status == 0, f"pstack {target} 执行错误：\n{output}")
            lines = output.splitlines(keepends=True)

        threads_str = []
        for line in lines:
            if line.strip() == '':
                continue
            if line.startswith('Thread '):
                threads_str.append('')
            threads_str[-1] += line

        res = ThreadList()
        for thread in threads_str:
            res.append(Thread(thread))

        return res


###############################
# 根据线程角色或者堆栈进行分类和统计
###############################
class GroupOperate(Operate):
    def __init__(self):
        super().__init__('group', 1, ThreadList, False, "")

    def execute(self, thread_list: ThreadList, params, reverse=False):
        Assert(params[0] in ["role", "bt", "ft"], f"错误：{params[0]}分组依据只支持 role、 bt、ft。")
        res = ThreadGroup(params[0])
        for thread in thread_list.threads:
            if params[0] == "role":
                res.append(thread.role.name, thread)
            elif params[0] == "bt":
                backtraces = "\n".join(thread.backtraces)
                res.append(backtraces, thread)
            elif params[0] == "ft":
                functraces = "\n".join(thread.functraces)
                res.append(functraces, thread)
        return res


###############################
# 打开某个分组
###############################
class OpenOperate(Operate):
    def __init__(self):
        super().__init__('open', 1, ThreadGroup, False, "")

    def execute(self, thread_group: ThreadGroup, params, reverse=False):
        Assert(params[0].isdigit(), f"错误：open操作后边需要接分组id。")
        return thread_group.open(int(params[0]))


###############################
# 根据线程角色过滤线程
###############################
class RoleOperate(Operate):
    def __init__(self):
        super().__init__('role', -1, ThreadList, True, "")

    def execute(self, thread_list: ThreadList, roles, reverse=False):
        res = ThreadList()
        for thread in thread_list.threads:
            if not reverse and thread.role.name in roles:
                res.append(thread)
            if reverse and thread.role.name not in roles:
                res.append(thread)

        return res


###############################
# 根据knl_thread_role的id，过滤线程
###############################
class RoleIDOperate(Operate):
    def __init__(self):
        super().__init__('roleid', -1, ThreadList, True, "")

    def execute(self, thread_list: ThreadList, roleids, reverse=False):
        res = ThreadList()
        for thread in thread_list.threads:
            if not reverse and str(thread.role.id) in roleids:
                res.append(thread)
            if reverse and str(thread.role.id) not in roleids:
                res.append(thread)

        return res


###############################
# 根据线程号，过滤线程
###############################
class ThreadIDOperate(Operate):
    def __init__(self):
        super().__init__('tid', -1, ThreadList, True, "")

    def execute(self, thread_list: ThreadList, tids, reverse=False):
        res = ThreadList()
        for thread in thread_list.threads:
            if not reverse and (str(thread.id) in tids or str(thread.tid) in tids):
                res.append(thread)
            if reverse and (str(thread.id) not in tids and str(thread.tid) not in tids):
                res.append(thread)

        return res


###############################
# 根据特定关键字的线程，过滤线程
###############################
class KeywordOperate(Operate):
    def __init__(self):
        super().__init__('key', -1, ThreadList, True, "")

    def execute(self, thread_list: ThreadList, keywords, reverse=False):
        res = ThreadList()
        for thread in thread_list.threads:
            for keyword in keywords:
                if thread.src.find(keyword) >= 0:
                    res.append(thread)
                    break
        if not reverse:
            return res

        v_res = ThreadList()
        for thread in thread_list.threads:
            if thread not in res.threads:
                v_res.append(thread)
        return v_res


class GrepOperate(KeywordOperate):
    def __init__(self):
        super().__init__()
        self.name = 'grep'


###############################
# 根据特定函数调用，过滤线程
###############################
class FuncOperate(Operate):
    def __init__(self):
        super().__init__('func', -1, ThreadList, True, "")

    def execute(self, thread_list: ThreadList, funcs, reverse=False):
        res = ThreadList()
        for thread in thread_list.threads:
            for func in funcs:
                if func in thread.functraces_func:
                    res.append(thread)
                    break
        if not reverse:
            return res

        v_res = ThreadList()
        for thread in thread_list.threads:
            if thread not in res.threads:
                v_res.append(thread)
        return v_res


###############################
# 根据堆栈深度，过滤线程
###############################
class DeepOperate(Operate):
    def __init__(self):
        super().__init__('deep', 1, ThreadList, True, "")

    def execute(self, thread_list: ThreadList, params, reverse=False):
        Assert(params[0].isdigit(), "错误：deep 需要一个数字")
        deep = int(params[0])
        res = ThreadList()
        for thread in thread_list.threads:
            if not reverse and len(thread.functraces) >= deep:
                res.append(thread)
            if reverse and len(thread.functraces) < deep:
                res.append(thread)
        return res


###############################
# 对线程的堆栈进行切片操作
###############################
class CutOperate(Operate):
    def __init__(self):
        super().__init__('cut', 1, ThreadList, False, "")

    def execute(self, thread_list: ThreadList, params, reverse=False):
        numbers = params[0].split(':')
        if len(numbers) == 1:
            num1 = ''
            num2 = numbers[0]
        else:
            num1 = numbers[0]
            num2 = numbers[1]
        start = eval(num1) if num1 != '' else None
        end = eval(num2) if num2 != '' else None

        for thread in thread_list.threads:
            thread.backtraces = thread.backtraces[start:end]
            thread.functraces = thread.functraces[start:end]
            thread.functraces_func = thread.functraces_func[start:end]
            thread.src = thread.title + '\n' + '\n'.join(thread.backtraces)

        return thread_list


#######################################################################################################################
# 执行器
#######################################################################################################################
class Executor(object):
    def __init__(self, argv):
        self.argv = argv
        self._operators = {}

    def register_operator(self, operator):
        self._operators[operator.name] = operator

    def done(self):
        return len(self.argv) == 0

    def exec_next_operator(self, analyze_res):
        if self.done():
            return analyze_res

        operator_name = self.argv.pop(0)
        operator = self._operators.get(operator_name)
        Assert(operator is not None, f"错误，未知的操作 {operator_name}")
        Assert(type(analyze_res) is operator.input_type,
               f"错误{operator_name}的输入需要是一个{operator.input_type.__name__}")

        reverse = False
        if operator.support_reverse and len(self.argv) > 0 and self.argv[0].lower() in ['v', '-v']:
            self.argv.pop(0)
            reverse = True

        params = []
        if operator.param_count == 0:
            pass
        elif operator.param_count > 0:
            Assert(len(self.argv) >= operator.param_count, f"错误，{operator_name}不足{operator.param_count}个参数。")
            params = self.argv[0:operator.param_count]
            self.argv = self.argv[operator.param_count:]
        else:
            while len(self.argv) > 0 and self.argv[0] not in self._operators:
                params.append(self.argv.pop(0))

        return operator.execute(analyze_res, params, reverse)


if __name__ == '__main__':
    if len(sys.argv) == 1 or sys.argv[1].lower() in ['-h', '--help', '-?']:
        print(f"""
介绍：
    一个分析pstack的工具，用来在pstack线程非常多的时候，快速对线程基于角色、堆栈进行分类，快速基于角色、函数、关键字进行线程的过滤，裁剪简化堆栈等。
    例如找到所有的worker、checkpointer、所有的含有某个函数调用、关键字的堆栈等。
    传统的pstack查看方法，例如 "cat pstack.log | grep xx | grep xxx ..." 是基于行的，并不能带来很好的观察效果。
    此工具针对这种使用场景进行设计，可以基于一个完整堆栈进行连续、多次的过滤、分组、统计等操作，更加有助于观察分析。

使用方法：
    python3 pstack-analyzer.py [PSTACK-FILE or PID] [operate1 param1, param2, ...] [operate2 [-v] param1 param2...] ...


参数：
    第一个参数默认是一个pstack的堆栈文件，如果没有这个文件则认为是一个进程号。
    后面的表示具体的操作、过滤或分组方式。
    
    操作        参数         支持-v    功能
    group      分组的依据     否       分组分类统计。依据可以是角色(role)、堆栈(bt)或函数堆栈(ft,忽略掉参数，仅看函数名的堆栈)
    open       分组的编号     否       分组之后，单独查看其中指定编号的组。
    
    role       角色名列表     是       找到所有这些角色的线程
    roleid     角色号列表     是       找到所有这些角色号的线程，角色号是 knl_thread_role 的数字，DMS的是10001，其他不知道的都是-1
    tid        线程号列表     是       找到所有这些线程号的线程，线程号可以指堆栈内的线程id，也可以指操作系统的LWP id
    key/grep   关键词列表     是       找到所有堆栈内包含这些关键字的线程
    func       函数名列表     是       找到所有存在这些函数调用的线程
    deep       大于等于几     是       找到所有堆栈深度大于等于指定深度的线程
    cut        切片的范围     否       裁切堆栈，裁切掉一部分堆栈以便观察。参数是python的切片语法, "结束", ":结束", "开始:", "开始:结束"

示例：
    1、  python3 pstack-analyzer.py pstackfile role worker func ExecLimit ExecAgg
        表示 找出所有的worker，继续找出所有调用了ExecLimit或ExecAgg的线程
        
    2、  python3 pstack-analyzer.py pstackfile key "postmaster.cpp" "snapshot=<optimized out>" func -v HeapamTupleUpdate ProcSleep
        表示 找出所有的含有关键字"postmaster.cpp"或"snapshot=<optimized out>"的，继续找出所有没有调用HeapamTupleUpdate和ProcSleep的线程
        
    3、  python3 pstack-analyzer.py pstackfile role worker func ExecLimit ExecAgg group bt
        表示 找出所有的worker，继续找出所有调用了ExecLimit或ExecAgg的线程，最后根据堆栈进行分类

注意：过滤时，指令参数的单词作为过滤参数时会有问题。例如 python3 pstack-analyzer.py file key role, 可以理解成找到所有含有
关键词'role'的线程，也可以理解成找到所有含关键词''的、角色为''的线程。

支持的角色与角色号列表：
    {ThrRole.get_role_map()}
    """)
        exit(0)

    pstack_src = sys.argv[1]
    operate_queue = sys.argv[2:]

    executor = Executor(operate_queue)
    executor.register_operator(GroupOperate())
    executor.register_operator(OpenOperate())
    executor.register_operator(RoleOperate())
    executor.register_operator(RoleIDOperate())
    executor.register_operator(ThreadIDOperate())
    executor.register_operator(KeywordOperate())
    executor.register_operator(GrepOperate())
    executor.register_operator(FuncOperate())
    executor.register_operator(DeepOperate())
    executor.register_operator(CutOperate())

    result = Pstack.parse(pstack_src)
    while not executor.done():
        result = executor.exec_next_operator(result)

    result.show()
